#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import asyncio
import json
import logging
import time
import uuid
import threading
from typing import Dict, Optional

import zmq
import zmq.asyncio
from aiortc import RTCPeerConnection, RTCSessionDescription, RTCIceCandidate
from aiortc import RTCDataChannel, VideoStreamTrack
from aiortc.contrib.media import MediaPlayer

from video_generator import VideoGenerator

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class RemoteServer:
    def __init__(self, server_id: str = None, signaling_host: str = "localhost", signaling_port: int = 5555):
        self.server_id = server_id or f"remote_server_{uuid.uuid4().hex[:8]}"
        self.signaling_host = signaling_host
        self.signaling_port = signaling_port
        
        # ZMQ配置
        self.zmq_context = None
        self.zmq_socket = None
        self.is_connected = False
        
        # WebRTC连接管理
        self.peer_connections: Dict[str, RTCPeerConnection] = {}
        self.data_channels: Dict[str, RTCDataChannel] = {}
        
        # 视频生成器
        self.video_generator = VideoGenerator()
        
        # 重连配置
        self.reconnect_attempts = 0
        self.max_reconnect_attempts = 5
        self.reconnect_delay = 5
        
        logger.info(f"远程服务器初始化: {self.server_id}")
    
    async def start(self):
        """启动远程服务器"""
        logger.info("启动远程服务器...")
        
        # 初始化ZMQ
        await self.init_zmq()
        
        # 启动视频生成器
        await self.video_generator.start()
        
        # 注册到信令服务器
        await self.register_to_signaling_server()
        
        # 启动消息处理循环
        await self.message_loop()
    
    async def init_zmq(self):
        """初始化ZMQ连接"""
        try:
            self.zmq_context = zmq.asyncio.Context()
            self.zmq_socket = self.zmq_context.socket(zmq.DEALER)
            self.zmq_socket.identity = self.server_id.encode('utf-8')
            
            zmq_url = f"tcp://{self.signaling_host}:{self.signaling_port}"
            self.zmq_socket.connect(zmq_url)
            
            logger.info(f"ZMQ连接到: {zmq_url}")
            self.is_connected = True
            self.reconnect_attempts = 0
            
        except Exception as e:
            logger.error(f"ZMQ连接失败: {e}")
            await self.handle_connection_error()
    
    async def register_to_signaling_server(self):
        """注册到信令服务器"""
        if not self.is_connected:
            return
            
        register_msg = {
            'type': 'register',
            'server_id': self.server_id,
            'timestamp': time.time()
        }
        
        try:
            await self.zmq_socket.send_string(json.dumps(register_msg))
            logger.info("已注册到信令服务器")
        except Exception as e:
            logger.error(f"注册失败: {e}")
            await self.handle_connection_error()
    
    async def message_loop(self):
        """消息处理主循环"""
        while True:
            try:
                if not self.is_connected:
                    await asyncio.sleep(1)
                    continue
                
                # 接收ZMQ消息
                try:
                    message = await asyncio.wait_for(
                        self.zmq_socket.recv_string(zmq.NOBLOCK), 
                        timeout=0.1
                    )
                    if message.strip():  # 确保消息不为空
                        try:
                            data = json.loads(message)
                            await self.handle_signaling_message(data)
                        except json.JSONDecodeError as e:
                            logger.error(f"JSON解析错误: {e}, 消息内容: {repr(message)}")
                    else:
                        logger.warning("收到空消息")
                except asyncio.TimeoutError:
                    pass
                except zmq.Again:
                    pass
                
                # 发送心跳
                await self.send_heartbeat()
                
                await asyncio.sleep(0.01)
                
            except Exception as e:
                logger.error(f"消息循环错误: {e}")
                await self.handle_connection_error()
                await asyncio.sleep(1)
    
    async def handle_signaling_message(self, data: dict):
        """处理信令消息"""
        try:
            message_type = data.get('type')
            source = data.get('source')
            
            logger.info(f"收到信令消息: {message_type} from {source}")
            
            if message_type == 'offer':
                await self.handle_offer(data)
            elif message_type == 'ice-candidate':
                await self.handle_ice_candidate(data)
            else:
                logger.warning(f"未知消息类型: {message_type}")
                
        except Exception as e:
            logger.error(f"处理信令消息错误: {e}")
    
    async def handle_offer(self, data: dict):
        """处理offer消息"""
        source = data.get('source')
        offer_data = data.get('offer')
        
        if not source or not offer_data:
            logger.error("无效的offer消息")
            return
        
        # 处理offer格式
        if isinstance(offer_data, dict):
            offer = offer_data
        else:
            logger.error(f"无效的offer格式: {offer_data}")
            return
        
        try:
            # 创建新的PeerConnection
            pc = RTCPeerConnection()
            self.peer_connections[source] = pc
            
            # 设置事件处理器
            await self.setup_peer_connection_handlers(pc, source)
            
            # 设置远程描述（必须在添加轨道之前）
            try:
                logger.info(f"设置远程描述: type={offer['type']}, sdp长度={len(offer['sdp'])}")
                logger.info(f"offer SDP内容: {offer['sdp'][:300]}...")  # 打印offer SDP前300字符用于调试
                await pc.setRemoteDescription(RTCSessionDescription(
                    sdp=offer['sdp'],
                    type=offer['type']
                ))
                logger.info("远程描述设置成功")
            except Exception as desc_error:
                logger.error(f"设置远程描述失败: {desc_error}")
                raise
            
            # 添加视频轨道
            try:
                video_track = self.video_generator.get_video_track()
                logger.info(f"获取视频轨道成功: {type(video_track)}")
                pc.addTrack(video_track)
                logger.info("视频轨道添加成功")
            except Exception as track_error:
                logger.error(f"添加视频轨道失败: {track_error}")
                raise
            
            # 设置轨道方向
            try:
                for transceiver in pc.getTransceivers():
                    if transceiver.sender.track and hasattr(transceiver.sender.track, 'kind'):
                        if transceiver.sender.track.kind == 'video':
                            # 远程服务器发送视频，所以设置为sendonly
                            transceiver.direction = 'sendonly'
                            logger.info(f"设置视频轨道方向: {transceiver.direction}")
            except Exception as direction_error:
                logger.error(f"设置轨道方向失败: {direction_error}")
            
            # 创建并发送answer
            try:
                logger.info("开始创建answer")
                # 确保所有transceiver都有正确的方向设置
                for transceiver in pc.getTransceivers():
                    if transceiver.sender.track and hasattr(transceiver.sender.track, 'kind'):
                        if transceiver.sender.track.kind == 'video':
                            # 确保视频轨道在answer中被包含，远程服务器发送视频
                            transceiver.direction = 'sendonly'
                            logger.info(f"确认视频轨道方向: {transceiver.direction}")
                    elif hasattr(transceiver, 'receiver') and transceiver.receiver.track:
                        # 对于接收轨道，设置为recvonly
                        if hasattr(transceiver.receiver.track, 'kind') and transceiver.receiver.track.kind == 'video':
                            transceiver.direction = 'recvonly'
                            logger.info(f"设置接收轨道方向: {transceiver.direction}")
                
                answer = await pc.createAnswer()
                logger.info(f"answer创建成功: type={answer.type}, sdp长度={len(answer.sdp)}")
                logger.info(f"answer SDP内容: {answer.sdp[:200]}...")  # 打印SDP前200字符用于调试
                
                logger.info("开始设置本地描述")
                await pc.setLocalDescription(answer)
                logger.info("本地描述设置成功")
            except Exception as answer_error:
                logger.error(f"创建answer失败: {answer_error}")
                import traceback
                logger.error(f"详细错误信息: {traceback.format_exc()}")
                raise
            
            response = {
                'type': 'answer',
                'answer': {
                    'sdp': pc.localDescription.sdp,
                    'type': pc.localDescription.type
                },
                'target': source
            }
            
            await self.send_signaling_message(response)
            logger.info(f"发送Answer到: {source}")
            
        except Exception as e:
            logger.error(f"处理offer错误: {e}")
    
    async def handle_ice_candidate(self, data: dict):
        """处理ICE候选"""
        source = data.get('source')
        candidate_data = data.get('candidate')
        
        if source not in self.peer_connections:
            logger.warning(f"未找到连接: {source}")
            return
        
        try:
            pc = self.peer_connections[source]
            
            # 处理前端发送的RTCIceCandidate对象格式
            if isinstance(candidate_data, dict):
                if 'candidate' in candidate_data and 'sdpMid' in candidate_data:
                    # 标准WebRTC格式 - 解析candidate字符串
                    candidate_str = candidate_data['candidate']
                    sdp_mid = candidate_data['sdpMid']
                    sdp_mline_index = candidate_data.get('sdpMLineIndex')
                    
                    # 解析candidate字符串
                    parts = candidate_str.split()
                    if len(parts) >= 8:
                        foundation = parts[0].split(':')[1]
                        component = int(parts[1])
                        protocol = parts[2]
                        priority = int(parts[3])
                        ip = parts[4]
                        port = int(parts[5])
                        typ = parts[7]
                        
                        candidate = RTCIceCandidate(
                            component=component,
                            foundation=foundation,
                            ip=ip,
                            port=port,
                            priority=priority,
                            protocol=protocol,
                            type=typ,
                            sdpMid=sdp_mid,
                            sdpMLineIndex=sdp_mline_index
                        )
                    else:
                        logger.error(f"无法解析ICE候选字符串: {candidate_str}")
                        return
                else:
                    # 自定义格式
                    candidate = RTCIceCandidate(
                        component=candidate_data.get('component', 1),
                        foundation=candidate_data.get('foundation', ''),
                        ip=candidate_data.get('address', candidate_data.get('ip', '')),
                        port=candidate_data.get('port', 0),
                        priority=candidate_data.get('priority', 0),
                        protocol=candidate_data.get('protocol', 'udp'),
                        type=candidate_data.get('type', 'host')
                    )
            else:
                logger.error(f"无效的ICE候选格式: {candidate_data}")
                return
            
            await pc.addIceCandidate(candidate)
            logger.info(f"添加ICE候选: {source}")
            
        except Exception as e:
            logger.error(f"处理ICE候选错误: {e}, 候选数据: {candidate_data}")
    
    async def setup_peer_connection_handlers(self, pc: RTCPeerConnection, client_id: str):
        """设置PeerConnection事件处理器"""
        
        @pc.on("icecandidate")
        async def on_icecandidate(candidate):
            if candidate:
                response = {
                    'type': 'ice-candidate',
                    'candidate': {
                        'component': candidate.component,
                        'foundation': candidate.foundation,
                        'address': candidate.ip,
                        'port': candidate.port,
                        'priority': candidate.priority,
                        'protocol': candidate.protocol,
                        'type': candidate.type
                    },
                    'target': client_id
                }
                await self.send_signaling_message(response)
        
        @pc.on("datachannel")
        def on_datachannel(channel):
            logger.info(f"数据通道建立: {client_id}")
            self.data_channels[client_id] = channel
            
            @channel.on("message")
            def on_message(message):
                logger.info(f"收到消息 from {client_id}: {message}")
                # 回复消息
                reply = f"收到<{message}>"
                try:
                    channel.send(reply)
                    logger.info(f"回复消息 to {client_id}: {reply}")
                except Exception as e:
                    logger.error(f"发送回复失败: {e}")
            
            @channel.on("close")
            def on_close():
                logger.info(f"数据通道关闭: {client_id}")
                if client_id in self.data_channels:
                    del self.data_channels[client_id]
        
        @pc.on("connectionstatechange")
        async def on_connectionstatechange():
            state = pc.connectionState
            logger.info(f"连接状态变化 {client_id}: {state}")
            
            if state == "failed" or state == "disconnected":
                logger.warning(f"连接断开: {client_id}")
                await self.cleanup_connection(client_id)
            elif state == "connected":
                logger.info(f"WebRTC连接建立成功: {client_id}")
    
    async def cleanup_connection(self, client_id: str):
        """清理连接"""
        if client_id in self.peer_connections:
            try:
                await self.peer_connections[client_id].close()
            except:
                pass
            del self.peer_connections[client_id]
        
        if client_id in self.data_channels:
            del self.data_channels[client_id]
        
        logger.info(f"清理连接: {client_id}")
    
    async def send_signaling_message(self, message: dict):
        """发送信令消息"""
        if not self.is_connected:
            logger.error("ZMQ未连接，无法发送消息")
            return
        
        try:
            await self.zmq_socket.send_string(json.dumps(message))
        except Exception as e:
            logger.error(f"发送信令消息失败: {e}")
            await self.handle_connection_error()
    
    async def send_heartbeat(self):
        """发送心跳"""
        # 每30秒发送一次心跳
        if hasattr(self, '_last_heartbeat'):
            if time.time() - self._last_heartbeat < 30:
                return
        
        heartbeat = {
            'type': 'heartbeat',
            'server_id': self.server_id,
            'timestamp': time.time(),
            'active_connections': len(self.peer_connections)
        }
        
        await self.send_signaling_message(heartbeat)
        self._last_heartbeat = time.time()
    
    async def handle_connection_error(self):
        """处理连接错误"""
        self.is_connected = False
        
        if self.reconnect_attempts < self.max_reconnect_attempts:
            self.reconnect_attempts += 1
            logger.warning(f"尝试重连 ({self.reconnect_attempts}/{self.max_reconnect_attempts})...")
            
            await asyncio.sleep(self.reconnect_delay * self.reconnect_attempts)
            
            try:
                if self.zmq_socket:
                    self.zmq_socket.close()
                await self.init_zmq()
                await self.register_to_signaling_server()
            except Exception as e:
                logger.error(f"重连失败: {e}")
        else:
            logger.error("重连次数超限，停止重连")
    
    async def stop(self):
        """停止服务器"""
        logger.info("停止远程服务器...")
        
        # 关闭所有WebRTC连接
        for client_id in list(self.peer_connections.keys()):
            await self.cleanup_connection(client_id)
        
        # 停止视频生成器
        await self.video_generator.stop()
        
        # 关闭ZMQ连接
        if self.zmq_socket:
            self.zmq_socket.close()
        
        if self.zmq_context:
            self.zmq_context.term()
        
        logger.info("远程服务器已停止")

async def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='WebRTC远程服务器')
    parser.add_argument('--server-id', default=None, help='服务器ID')
    parser.add_argument('--signaling-host', default='localhost', help='信令服务器地址')
    parser.add_argument('--signaling-port', type=int, default=5555, help='信令服务器端口')
    
    args = parser.parse_args()
    
    server = RemoteServer(
        server_id=args.server_id,
        signaling_host=args.signaling_host,
        signaling_port=args.signaling_port
    )
    
    try:
        await server.start()
    except KeyboardInterrupt:
        logger.info("收到中断信号")
    finally:
        await server.stop()

if __name__ == '__main__':
    asyncio.run(main())